#include "CNN/mpi_mnist_base.hpp"
#include "intel_itt_wrapper.hpp"
// #include <gtest/gtest.h>

// Where to find the MNIST dataset.
std::string kDataRoot = "build/data";
// Defualt BGV Params Set
params<helib::BGV> mpi_bgv_params;

auto main(int argc, char **argv) -> int {
  // INTEL_ITT_PAUSE; //! ITT - Stop
  MPI_Init(&argc, &argv);
  // int provided;
  // MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided);
  // if (provided < MPI_THREAD_MULTIPLE) {
  //   printf("The threading support level is lesser than that demanded.\n");
  //   MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
  // } else {
  //   printf("The threading support level corresponds to that demanded.\n");
  // }
  // int ret = MPI_Init(&argc, &argv);
  // EXPECT_EQ(ret, 0);
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  // ::testing::TestEventListeners &listeners =
  //     ::testing::UnitTest::GetInstance()->listeners();
  // if (rank != COMM_WORLD_ROOT) {
  //   //* Only the root has a secret key
  //   delete listeners.Release(listeners.default_result_printer());
  // }
  torch::manual_seed(1);
  unsigned long test_num = 1;

  // Create a `command_line_parser` using `params<helib::BGV>`
  command_line_parser parser(mpi_bgv_params);
  parser.add_option("-tn", "2", "test_num", &test_num);

  //* Parse command line
  int ret = parser.parse(argc, argv, rank == COMM_WORLD_ROOT);
  // EXPECT_EQ(ret, 0);

  torch::Device device(torch::kCPU);
  Net model(COMM_WORLD_ROOT, MPI_COMM_WORLD, mpi_bgv_params);
  model.to(device);

  std::string pt_file = kDataRoot + "/emnist.pt";
  load_model(model, device, pt_file);

  test(model, device, kDataRoot, test_num);
  MPI_Finalize();
}
